import onnx
import onnxruntime as ort
import numpy as np
from datetime import datetime
import json
from numpy.linalg import norm

def add_intermediate_outputs(model_path, save_path=None):
    """
    修改ONNX模型，将所有中间层节点添加为输出
    """
    if save_path is None:
        save_path = model_path.replace('.onnx', '_with_intermediate.onnx')
    
    model = onnx.load(model_path)
    intermediate_outputs = []
    
    for node in model.graph.node:
        for output in node.output:
            if output not in [o.name for o in model.graph.output]:
                value_info = onnx.helper.ValueInfoProto()
                value_info.name = output
                intermediate_outputs.append(value_info)
    
    model.graph.output.extend(intermediate_outputs)
    onnx.save(model, save_path)
    return save_path

def get_model_outputs(model_path, input_data):
    """
    获取模型所有层的输出
    """
    session = ort.InferenceSession(model_path)
    input_name = session.get_inputs()[0].name
    output_names = [output.name for output in session.get_outputs()]
    outputs = session.run(output_names, {input_name: input_data})
    
    return dict(zip(output_names, outputs))

def calculate_cosine_similarity(tensor1, tensor2):
    """
    计算两个张量之间的余弦相似度，增加数值稳定性
    """
    # 将张量展平
    flat1 = tensor1.flatten()
    flat2 = tensor2.flatten()
    
    # 检查是否全为零
    if np.all(flat1 == 0) or np.all(flat2 == 0):
        return 0.0
    
    try:
        # 使用numpy的norm计算余弦相似度
        cos_sim = np.dot(flat1, flat2) / (norm(flat1) * norm(flat2))
        # 处理数值误差导致的越界
        cos_sim = np.clip(cos_sim, -1.0, 1.0)
        return float(cos_sim)
    except Exception as e:
        print(f"计算余弦相似度时出错: {e}")
        print(f"tensor1 统计: mean={np.mean(flat1)}, std={np.std(flat1)}, min={np.min(flat1)}, max={np.max(flat1)}")
        print(f"tensor2 统计: mean={np.mean(flat2)}, std={np.std(flat2)}, min={np.min(flat2)}, max={np.max(flat2)}")
        return 0.0

def compare_models_outputs(fp32_outputs, int8_outputs):
    """
    比较两个模型的输出并计算相似度
    """
    similarities = {}
    common_layers = set(fp32_outputs.keys()) & set(int8_outputs.keys())
    
    for layer_name in common_layers:
        fp32_tensor = fp32_outputs[layer_name]
        int8_tensor = int8_outputs[layer_name]
        
        # 检查形状是否匹配
        if fp32_tensor.shape != int8_tensor.shape:
            print(f"警告: 层 {layer_name} 的形状不匹配 - FP32: {fp32_tensor.shape}, INT8: {int8_tensor.shape}")
            continue
        
        # 计算相似度
        similarity = calculate_cosine_similarity(fp32_tensor, int8_tensor)
        
        # 保存结果
        similarities[layer_name] = {
            'cosine_similarity': float(similarity),
            'shape': fp32_tensor.shape,
            'fp32_stats': {
                'mean': float(np.mean(fp32_tensor)),
                'std': float(np.std(fp32_tensor)),
                'min': float(np.min(fp32_tensor)),
                'max': float(np.max(fp32_tensor))
            },
            'int8_stats': {
                'mean': float(np.mean(int8_tensor)),
                'std': float(np.std(int8_tensor)),
                'min': float(np.min(int8_tensor)),
                'max': float(np.max(int8_tensor))
            }
        }
    
    return similarities

def save_results(results, base_path):
    """
    保存比较结果到JSON文件
    """
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = f"{base_path}/similarity_comparison_{timestamp}.json"
    
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, indent=4, ensure_ascii=False)
    
    print(f"\n结果已保存到文件: {output_file}")
    return output_file

def print_results(results):
    """
    打印比较结果
    """
    print("\n模型量化前后各层余弦相似度比较:")
    print("-" * 100)
    print(f"{'层名称':<50} {'相似度':>10} {'形状':<20}")
    print("-" * 100)
    
    # 按相似度排序（从低到高）
    sorted_results = sorted(results.items(), key=lambda x: x[1]['cosine_similarity'])
    
    for layer_name, info in sorted_results:
        similarity = info['cosine_similarity']
        shape = str(info['shape'])
        print(f"{layer_name:<50} {similarity:>10.6f} {shape:<20}")
        
        # 如果相似度较低，打印详细统计信息
        if similarity < 0.95:
            print(f"  FP32统计: mean={info['fp32_stats']['mean']:.6f}, std={info['fp32_stats']['std']:.6f}, "
                  f"min={info['fp32_stats']['min']:.6f}, max={info['fp32_stats']['max']:.6f}")
            print(f"  INT8统计: mean={info['int8_stats']['mean']:.6f}, std={info['int8_stats']['std']:.6f}, "
                  f"min={info['int8_stats']['min']:.6f}, max={info['int8_stats']['max']:.6f}")
            print("-" * 100)

if __name__ == "__main__":
    try:
        # 模型路径
        fp32_model = "your_model.onnx"
        int8_model = "your_model_quant.onnx"
        
        # 修改模型以输出中间层
        print("正在修改FP32模型...")
        fp32_modified = add_intermediate_outputs(fp32_model)
        print("正在修改INT8模型...")
        int8_modified = add_intermediate_outputs(int8_model)
        
        # 创建输入数据
        print("准备输入数据...")
        input_data = np.random.randn(1, 3, 640, 640).astype(np.float32)
        
        # 获取两个模型的输出
        print("获取FP32模型输出...")
        fp32_outputs = get_model_outputs(fp32_modified, input_data)
        print("获取INT8模型输出...")
        int8_outputs = get_model_outputs(int8_modified, input_data)
        
        # 比较输出并计算相似度
        print("计算相似度...")
        results = compare_models_outputs(fp32_outputs, int8_outputs)
        
        # 保存结果
        save_path = "your_save_path"
        output_file = save_results(results, save_path)
        
        # 打印结果
        print_results(results)
        
    except Exception as e:
        print(f"发生错误: {e}")
        import traceback
        traceback.print_exc()